import numpy as np
import logging
from crowd_sim.envs.policy.policy import Policy
from crowd_sim.envs.utils.action import ActionXY

class From_data(Policy):
    def __init__(self):
        super().__init__()
        self.name = 'data'
        self.trainable = False
        self.multiagent_training = True
        self.kinematics = 'holonomic'

    def configure(self, config):
        return

    def set_phase(self, phase):
        return

    def predict(self, state):
        """
         I have replaced predict with read_data. But in the base class: Policy, predict is a abstract method.
        So I have to do this, but this function can not be called.
        """
        logging.warning("The predict function in From_data policy should not be called")
        return state

    def read_data(self, id, trajectories): # dict of dict of tuple: { frame_id: {pedestrian_id:(x,y), ......
        """
        
        """
        current_frame = min(trajectories.keys())
        current_x = trajectories[current_frame][id][0]
        current_y = trajectories[current_frame][id][1]
        next_x = trajectories[current_frame+1][id][0]
        next_y = trajectories[current_frame+1][id][1]
        v_x = (next_x - current_x)/self.time_step
        v_y = (next_y - current_y)/self.time_step
        action = ActionXY(v_x,v_y)
        return action

class CentralizedFrom_data(From_data):
    def __init__(self):
        super().__init__()

    def read_data(self, trajectories, current_pedestrian):  # dict of dict of tuple: { frame_id: {pedestrian_id:(x,y), ......
        """
        
        """
        current_frame = min(trajectories.keys())
        actions = []
        is_end = []
        for id in current_pedestrian:
            try:
                # print(current_frame)
                if id in trajectories[current_frame + 10.0].keys():  # 保证下一帧这个人还在
                    current_x = trajectories[current_frame][id][0]
                    current_y = trajectories[current_frame][id][1]
                    next_x = trajectories[current_frame + 10.0][id][0]
                    next_y = trajectories[current_frame + 10.0][id][1]
                    v_x = (next_x - current_x) / self.time_step
                    v_y = (next_y - current_y) / self.time_step
                    action = ActionXY(v_x, v_y)
                    actions.append(action)
                else:
                    # actions.append('End of trajectory')
                    is_end.append(id)
            except KeyError: # 把数据集跑完了
                exit()
        trajectories.pop(current_frame)
        return is_end, actions

